import cv2
import numpy as np
import os

def calculate_similarity(img1, img2):
    # 使用均方差计算两个图像之间的相似度
    difference = cv2.absdiff(img1, img2)
    similarity = np.sum(difference ** 2)
    return similarity

def extract_frames_and_calculate_similarity(video_path, output_dir):
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print("Error: Could not open video.")
        return

    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frame_similarity = []

    # 读取第一帧
    ret, prev_frame = cap.read()
    if not ret:
        print("Error: Could not read the first frame.")
        return

    # 保存第一帧
    frame_number = 0
    cv2.imwrite(os.path.join(output_dir, f"frame_{frame_number:04d}.jpg"), prev_frame)

    for i in range(1, frame_count):
        ret, current_frame = cap.read()
        if not ret:
            break

        # 计算当前帧与前一帧的相似度
        similarity = calculate_similarity(prev_frame, current_frame)
        frame_similarity.append(similarity)

        # 保存当前帧
        frame_number += 1
        cv2.imwrite(os.path.join(output_dir, f"frame_{frame_number:04d}.jpg"), current_frame)

        # 更新前一帧
        prev_frame = current_frame

    cap.release()

    # 将相似度数组保存到文件
    similarity_file_path = os.path.join(output_dir, "frame_similarity.npy")
    np.save(similarity_file_path, frame_similarity)

    print(f"Frames extracted and similarity calculated. Similarity array saved to {similarity_file_path}")

# 使用示例
video_path = 'az_recorder_20240712_125704.mp4'
output_dir = 'output_frames'
extract_frames_and_calculate_similarity(video_path, output_dir)
